package com.lizhi.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.lizhi.config.LocalTaskProperties;
import com.lizhi.entity.LocalTask;
import com.lizhi.entity.SubTaskExecuteInfo;
import com.lizhi.enums.SubTaskState;
import com.lizhi.enums.TaskTotalState;
import com.lizhi.exception.subex.*;
import com.lizhi.mapper.LocalTaskMapper;
import com.lizhi.service.TaskService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;

import java.util.HashSet;
import java.util.List;
import java.util.Objects;
import java.util.Set;

import static com.lizhi.enums.TaskTotalState.*;

@Slf4j
public class TaskServiceImpl implements TaskService {
    @Autowired
    private LocalTaskMapper localTaskMapper;

    @Autowired
    private LocalTaskProperties localTaskProperties;

    @Override
    public boolean publishTask(String taskType, List<String> businessKeys) {
        // 任务类型不得为空
        if (taskType == null || taskType.isBlank()) {
            throw new TaskTypeIsEmptyException();
        }
        LocalTask localTask = new LocalTask();
        localTask.setTaskType(taskType);
        int length = businessKeys.size();
        localTask.setBusinessKeys(businessKeys);
        localTask.setTaskTotalState(NOT_START);
        // 根据业务主键长度初始化子任务执行信息
        localTask.setSubTaskExecuteInfos(SubTaskExecuteInfo.initSubTaskExecuteInfos(length));
        // 根据业务主键长度初始化子任务状态
        localTask.setSubTaskStates(SubTaskState.initSubTaskState(length));
        return localTaskMapper.insert(localTask) > 0;
    }

    @Override
    public SubTaskState getSubTaskState(Integer taskId, Integer subTaskIndex) {
        return checkTaskIsExist(taskId, subTaskIndex).getSubTaskStates().get(subTaskIndex);
    }

    @Override
    public LocalTask getTaskById(Integer taskId) {
        LocalTask localTask = localTaskMapper.selectById(taskId);
        if (localTask == null) {
            throw new TaskNotFoundException();
        }
        return localTask;
    }

    /**
     * 验证主任务是否存在和子任务是否越界
     * 如果主任务存在并且子任务索引没有越界，返回主任务
     *
     * @param taskId
     * @return
     */
    private LocalTask checkTaskIsExist(Integer taskId, int subTaskIndex) {
        LocalTask localTask = localTaskMapper.selectById(taskId);
        if (localTask == null) {
            throw new TaskNotFoundException();
        }
        if (subTaskIndex >= localTask.getSubTaskStates().size()) {
            throw new SubTaskNotFoundException();
        }
        return localTask;
    }

    /**
     * 只有任务之前的状态是EXECUTING才允许变成FAILED或者SUCCESS
     *
     * @param subTaskState
     */
    private void checkSubTaskStateMustBeExecuting(SubTaskState subTaskState) {
        if (subTaskState != SubTaskState.EXECUTING) {
            throw new StateMustBeExecutingBeforeCompleteException();
        }
    }

    /**
     * 根据用户配置
     * 检查子任务是否连续
     *
     * @param subTaskStates
     * @param subTaskIndex
     */
    private void checkSubTaskContinuous(List<SubTaskState> subTaskStates, Integer subTaskIndex) {
        // 需要任务连续执行，所以不能跳着完成任务
        if (localTaskProperties.getContinuousTask()) {
            for (int i = 0; i < subTaskIndex; i++) {
                // subTaskIndex之前的任务存在没成功的
                if (SubTaskState.SUCCESS != subTaskStates.get(i)) {
                    throw new SubTaskNotContinuousException();
                }
            }
        }
    }

    @Override
    public boolean completeSubTask(Integer taskId, Integer subTaskIndex, String errorMessage) {
        // 先检查主任务和子任务
        LocalTask localTask = checkTaskIsExist(taskId, subTaskIndex);
        List<SubTaskState> subTaskStates = localTask.getSubTaskStates();
        // 检查子任务是否连续
        checkSubTaskContinuous(subTaskStates, subTaskIndex);
        // 完成子任务的时候，当前状态必须是执行中
        checkSubTaskStateMustBeExecuting(subTaskStates.get(subTaskIndex));
        int subTaskSize = subTaskStates.size();
        List<SubTaskExecuteInfo> subTaskExecuteInfos = localTask.getSubTaskExecuteInfos();
        SubTaskExecuteInfo subTaskExecuteInfo = subTaskExecuteInfos.get(subTaskIndex);
        // 子任务失败
        if (errorMessage != null && !errorMessage.isEmpty()) {
            // 获取子任务总失败次数
            Integer totalFailCount = subTaskExecuteInfos.stream().map(SubTaskExecuteInfo::getFailCount).reduce(0, Integer::sum);
            // 子任务失败
            subTaskStates.set(subTaskIndex, SubTaskState.FAILED);
            // 子任务总失败次数超过上限
            if (totalFailCount > localTaskProperties.getTaskMaxFailCount()) {
                // 设置主任务失败，后续再次调度的时候再次发邮件通知
                // 因为要保证任务一定要被消费，所以不能不处理FAILED的任务
                // TODO: 在后续被调度的时候发邮件通知
                localTask.setTaskTotalState(TaskTotalState.FAILED_OVER);
                log.error("任务[{}] 的子任务总失败次数超过 {} 次，请排查", taskId, totalFailCount);
            } else if (Objects.equals(subTaskExecuteInfo.getFailCount() + 1, localTaskProperties.getSubTaskMaxFailCount())) {
                // 单个子任务失败次数达到上限
                log.error("任务[{}] 的子任务失败次数超过 {} 次，请注意", taskId, localTaskProperties.getSubTaskMaxFailCount());
            } else {
                // 设置主任务状态为子任务失败
                localTask.setTaskTotalState(TaskTotalState.SUB_TASK_FAILED);
            }
            // 记录子任务失败原因
            subTaskExecuteInfo.fail(errorMessage);
        } else {
            // 子任务成功
            subTaskStates.set(subTaskIndex, SubTaskState.SUCCESS);
            // 子任务执行成功
            subTaskExecuteInfo.success();
            // 如果最后一个子任务完成，设置总任务状态
            if (subTaskIndex == subTaskSize - 1) {
                // TODO: 删除这条任务记录，将其添加到history表中
                localTask.setTaskTotalState(TaskTotalState.SUCCESS);
            }
        }
        // 更新主任务
        return localTaskMapper.updateById(localTask) >= 1;
    }

    @Override
    public boolean completeSubTask(Integer taskId, String errorMessage) {
        return completeSubTask(
                taskId,
                localTaskMapper
                        .selectOne(new LambdaQueryWrapper<LocalTask>().eq(LocalTask::getId, taskId))
                        // nextSubTaskIndex是下一个执行的子任务索引，所以正在执行的任务索引是nextSubTaskIndex-1
                        .getNextSubTaskIndex() - 1,
                errorMessage
        );
    }

    @Override
    public boolean startSubTask(Integer taskId, Integer subTaskIndex) {
        LocalTask localTask = checkTaskIsExist(taskId, subTaskIndex);
        List<SubTaskState> subTaskStates = localTask.getSubTaskStates();
        SubTaskState prevSubTaskState = subTaskStates.get(subTaskIndex);
        checkSubTaskContinuous(subTaskStates, subTaskIndex);
        // 连续执行、subTaskIndex是0-》√
        // 连续执行、subTaskIndex不为0-》×（checkSubTaskContinuous检查过了已经抛出异常了）
        // 不连续执行、subIndex为0-》√
        // 不连续执行、subIndex不为0-》√
        // 检查当前子任务是否已经执行成功
        checkSubTaskHashSuccess(prevSubTaskState);
        // 主任务状态变为执行中
        // 因为有可能是失败了的任务重新被下发执行了
        localTask.setTaskTotalState(TaskTotalState.EXECUTING);
        subTaskStates.set(subTaskIndex, SubTaskState.EXECUTING);
        // 更新下一个执行的子任务索引
        localTask.setNextSubTaskIndex(subTaskIndex + 1);
        return localTaskMapper.update(localTask,
                new LambdaQueryWrapper<LocalTask>()
                        .eq(LocalTask::getId, taskId)
                        // 乐观锁
                        // 如果同时有多个线程开启子任务，只能有一个得到成功的响应
                        .apply(String.format("JSON_EXTRACT(SUB_TASK_STATES,'$[%d]') = {0}", subTaskIndex), prevSubTaskState.getCode())) == 1;
    }

    /**
     * 已经执行成功的子任务不得重新执行
     *
     * @param subTaskState
     */
    private void checkSubTaskHashSuccess(SubTaskState subTaskState) {
        if (subTaskState == SubTaskState.SUCCESS) {
            throw new SubTaskHasSuccessException();
        }
    }

    @Override
    public boolean startSubTask(Integer taskId) {
        return startSubTask(
                taskId,
                localTaskMapper
                        .selectOne(new LambdaQueryWrapper<LocalTask>().eq(LocalTask::getId, taskId))
                        // 直接拿下一次子任务的索引
                        .getNextSubTaskIndex()
        );
    }

    @Override
    public List<LocalTask> dispatchTasks(String taskType, Integer shardSize, Integer shardIndex, Integer limit) {
        List<LocalTask> records = localTaskMapper.selectPage(
                new Page<>(shardIndex, limit),
                new LambdaQueryWrapper<LocalTask>()
                        .in(LocalTask::getTaskTotalState, NOT_START, SUB_TASK_FAILED, FAILED_OVER)
                        .eq(LocalTask::getTaskType, taskType)
                        // TODO: 这里等会试一下为什么删掉
//                        .apply(String.format("id %% %d = %d", shardSize, shardIndex))).getRecords();
                        .apply("id % {0} = {1}", shardSize, shardIndex)).getRecords();
        Set<Integer> removeIds = new HashSet<>();
        records.forEach(r -> {
            TaskTotalState prevTaskTotalState = r.getTaskTotalState();
            r.setTaskTotalState(DISPATCHED);
            // 乐观锁更新看看任务状态能不能更新成功
            int success = localTaskMapper.update(r, new LambdaQueryWrapper<LocalTask>().eq(LocalTask::getTaskTotalState, prevTaskTotalState).eq(LocalTask::getId, r.getId()));
            if (success <= 0) {
                removeIds.add(r.getId());
            }
        });
        for (LocalTask record : records) {
            // 这个任务取失败了
            if (removeIds.contains(record.getId())) {
                continue;
            }
            // TODO: 发邮件提醒
            if (record.getTaskTotalState() == FAILED_OVER) {
                Integer totalSubTaskFailCount = record.getSubTaskExecuteInfos().stream().map(SubTaskExecuteInfo::getFailCount).reduce(0, Integer::sum);
                log.error("任务[{}] 的子任务失败总数过多，第 {} 次尝试", record.getId(), totalSubTaskFailCount);
            }
            if (record.getTaskTotalState() == SUB_TASK_FAILED) {
                Integer subTaskFailCount = record.getSubTaskExecuteInfos().get(record.getNextSubTaskIndex() - 1).getFailCount();
                log.error("任务[{}] 的子任务失败次数过多，第 {} 次尝试", record.getId(), subTaskFailCount);
            }
        }
        return records;
    }

    @Override
    public List<LocalTask> dispatchTasks(String taskType, Integer limit) {
        return dispatchTasks(taskType, 1, 0, limit);
    }
}
